import numpy as np
from scipy.integrate import odeint
from pdb import set_trace
"""
Module to calculate the transient deformations, displacements
and forces from a Hertzian impact between a sphere and a load
cell.
"""
# Global constants
g = 9.8             # gravity (m/s^2)

class Sim():
    """
    Class to calculate the dynamics (Force, displacement) for a impact between
    a sphere and elastic half-space assuming a Hertzian dashpot model.
    """
    def __init__(self):
        self.dynamics = np.zeros((1), dtype=[('time','<f4'), ('deformation','<f4'),
            ('deformation rate', '<f4'), ('force', '<f4')])
        self.params = {
                'h':None,         # drop height (m)
                'R':None,        # sphere radius (m)
                'Es':None,         # sphere modulus (Pa)
                'Ep':None,         # plane modulus (Pa)
                'vs':None,          # sphere Poisson's ratio
                'vp':None,          # plane Poisson's ratio
                'm':None,       # sphere mass (kg)
                'COR':None           # coefficient of restitution
                }

    def set_default_params(self):
        """
        Set default parameter values for debugging
        """
        default_params = {
                'h':4.7e-2,       
                'R':4.76e-3,      
                'Es':3.0e9,       
                'Ep':3.0e9,       
                'vs':0.33,        
                'vp':0.33,        
                'm':0.407e-3,     
                'COR':0.8         
                }
        self.set_parameters(default_params)
        self.dynamics = np.zeros((10**5), dtype=[('time','<f4'), ('deformation','<f4'),
            ('deformation rate', '<f4'), ('force', '<f4')])
        self.dynamics['time'] = np.linspace(0,10**-3,10**5)
        
    def set_parameters(self,parameters):
        """
        Set the tablet and set-up parameters (all in SI units). 
        See self.params for dictionary.
        """
        if set(parameters.keys()).issubset(set(self.params.keys())):
            self.params.update(parameters)
        else: 
            raise ValueError('Invalid parameter in dictionary')
         
    def Hertz_constant(self,E1,E2,v1,v2,R):
        """
        function to calculate the effective stiffness used in the Hertz contact
        theory.
        E1 - tablet modulus (Pa)
        E2 - plane modulus (Pa)
        v1 - tablet Poisson ratio
        v2 - Plane Poisson ratio
        R - radius of sphere (m)
        return K - constant for Hertz force F = Ky^1.5
        """
        E_eff = 1. / ( (1.-v1**2)/E1 + (1.-v2**2)/E2 )
        K = 4./3.*R**0.5*E_eff
        return K  

    def Hertz_force(self,y,K):
        """
        function to calculate the force on the load cell
        y - indentation depth (m)
        K - Hertz constant (kg/(m^0.5*s^2))
        return F - total contact force between tablet and load cell (N)
        """
        if np.isscalar(y):
            y = np.max([y,0])
        else:
            y[y<0] = 0
        F = K*y**(1.5)
        return F
         
    def dissipative_constant(self,m,K,COR):
        """
        function to calculate the constant for the dissipative force 
        (see Tsuji Powder Tech, 71 (1992) 239)
        m - tablet mass (kg)
        K - Hertz constant (kg/(m^0.5*s^2))
        COR - coefficient of restitution for the tablet
        return h - constant for the dissipative force F = h*y**0.25*y_prime
        """
        a = -1*(np.sqrt(5)*np.log(COR)) / np.sqrt(np.log(COR)**2 + np.pi**2)
        h = a*np.sqrt(m*K)
        return h

    def dissipative_force(self,eta,y,y_prime):
        """
        function to calculate dissipative force (see Tsuji Powder Tech, 71
        (1992) 239)
        eta - scalar dissipative factor
        y - indentation depth (m)
        y_prime - velocity (rate of change of indentation depth) (m/s)
        return F - dissipation force
        """
        if np.isscalar(y):
            y = np.max([y,0])
        else:
            y[y<0] = 0
        F = eta * y**0.25 * y_prime
        return F

    def derivatives(self,y,t,m,K,eta):
        """
        function to calculate the derivatives
        y - dependent variable of coupled ODEs - 
        [indentation depth, rate change indentation depth] - [m,m/s]
        t - discrete time array for numerical integration (s)
        K -Hertz_constant (kg/(m^0.5*s^2))
        eta - dissipative constant kg*m^(1/4)/s
        return [2x1] numpy array derivatives
        """
        try:
            y0_prime = y[1]
            y1_prime = -1./m * (self.Hertz_force(y[0],K) 
                       + self.dissipative_force(eta,y[0],y[1]) - m*9.8)
            derivative = np.array([y0_prime,y1_prime])
        except (SystemExit, KeyboardInterrupt):
            raise
        return derivative 
        
    @classmethod
    def simulate(cls, t, m, R, h, Es, Ep, vs, vp, COR):
        """
        Calculate the deformation, deformation rate and force
        at time points given in t for passed parameters.
        t - array of numpy time points (s)
        h - drop height (m)
        R - sphere radius (m)
        Es - sphere modulus (Pa)
        Ep - plane modulus (Pa)
        vs - sphere Poisson's ratio
        vp - plane Poisson's ratio
        m - sphere mass (kg)
        COR - coefficient of restitution
        """
        sim = cls()
        self.set_parameters({'m':m, 'R':R, 'h':h, 'Es':Es, 'Ep':Ep, 
            'vs':vs, 'vp':vp, 'COR':COR})
        self.dynamics = np.zeros((t.shape[0]), dtype=[('time','<f4'), 
            ('deformation','<f4'), ('deformation rate', '<f4'), ('force', '<f4')])
        K = sim.Hertz_constant(Es,Ep,vs,vp,R)
        eta = sim.dissipative_constant(m,K,COR)
        vo = g*(2*h/g)**0.5
        y_init = [0,vo]
        y = odeint(sim.derivatives,y_init,t,(m,K,eta))
        mask = y[:,0] > 0
        #t = t[mask]
        #y = y[mask]
        #pts = np.sum(mask)
        #t = np.array([t[i] for i in range(0,pts,100)])
        #y = np.array([y[i,:] for i in range(0,pts,100)])
        f = sim.Hertz_force(y[mask][:,0],K) 
        mask = ~np.isnan(f)
        sim.dynamics['time'] = t #t[mask]
        sim.dynamics['deformation'] = y[0] #y[mask][0]
        sim.dynamics['deformation rate'] = y[1] #y[mask][1]
        sim.dynamics['force'] = f #f[mask]
        return sim

    def solve(self):
        """
        numerically integrate the differential equation
        return [[nx1],[nxk],[nx1]] [time,displacement,force]
            [[time (s)], 
            [deformation (m), rate of deformation (m/s)],
            [force (N)]]
        """
        [m, R, h, Es, Ep, vs, vp, COR] = [self.params[k] for k in 
                ['m', 'R', 'h', 'Es', 'Ep', 'vs', 'vp', 'COR']]
        K = self.Hertz_constant(Es,Ep,vs,vp,R)
        eta = self.dissipative_constant(m,K,COR)
        vo = g*(2*h/g)**0.5
        y_init = [0,vo]
        t = np.linspace(0,10**-3,10**5)
        y = odeint(self.derivatives,y_init,t,(m,K,eta))
        mask = y[:,0] > 0
        pts = np.sum(mask)
        t = t[mask]
        y = y[mask]
        t = np.array([t[i] for i in range(0,pts,100)])
        y = np.array([y[i,:] for i in range(0,pts,100)])
        f = self.Hertz_force(y[:,0],K) 
        mask = ~np.isnan(f)
        t = t[mask]
        y = y[mask]
        f = f[mask]
        return [t,y,f]

    def fit_material_props(self):
        """
        Fit the dynamic load cell voltage signal with the Hertzian and deflectable 
        crystal equations to extract the compact modulus, Ep, and Poisson ratio, vp.
        """
        sim = ResSim()
        t_exp, v_exp = self.crop_data(self.raw_data[0].copy(), 
                       self.raw_data[1].copy())
        params = dict(self.params.items() + self.estimates.items())
        d_exp = v_exp * params['cal']
        t_exp -= t_exp[np.argmax(d_exp)]

        def objective(p,t_exp,d_exp,params):
            params.update({'Ep':p[0], 'vp':p[1], 'COR':p[2]})
            sim.set_parameters({i:params[i] for i in ['h','R','Es','Ep','vs',
                'vp','m', 'COR','m_crystal','k_crystal','mu_dashpot']})
            [t,y,f] = sim.solve()
            t_sim, d_sim = self.crop_data(t,y[:,2])
            t_sim -= t_sim[np.argmax(d_sim)]
            spline = splrep(t_sim,d_sim)
            d_sim = splev(t_exp,spline)
            residuals = np.sum((d_exp - d_sim)**2)
            return residuals

        func_args = (t_exp, d_exp, params) 
        p_guess = np.array([self.estimates['Ep'], self.estimates['vp'], 
                            self.estimates['COR']])
        p_opt = fmin(objective, p_guess, args=func_args, xtol=10, ftol=10, 
                     maxiter=200, maxfun=50, full_output=1, disp=0, retall=0)

        self.estimates.update({'Ep':p_opt[0][0], 'vp':p_opt[0][1], 
                               'COR':p_opt[0][2]})

        SSres = p_opt[1]
        SStot = np.sum((d_exp - np.mean(d_exp))**2)
        R2 = 1 - SSres / SStot
        self.estimates.update({'R2' : R2})

    def fit_piezo_parameters(self):
        """
        Fit the voltage by varying the load cell parameters
        """
        sim = ResSim()
        t_exp, v_exp = self.crop_data(self.raw_data[0].copy(), 
                       self.raw_data[1].copy())
        params = dict(self.params.items() + self.estimates.items())
        d_exp = v_exp * params['cal']
        t_exp -= t_exp[np.argmax(d_exp)]

        def objective(p,t_exp,d_exp,params):
            params.update({'m_crystal':p[0], 'k_crystal':p[1], 'mu_dashpot':p[2]})
            sim.set_parameters({i:params[i] for i in ['h','R','Es','Ep','vs',
                'vp','m', 'COR','m_crystal','k_crystal','mu_dashpot']})
            [t,y,f] = sim.solve()
            t_sim, d_sim = self.crop_data(t,y[:,2])
            t_sim -= t_sim[np.argmax(d_sim)]
            spline = splrep(t_sim,d_sim)
            d_sim = splev(t_exp,spline)
            residuals = np.sum((d_exp - d_sim)**2)
            return residuals

        func_args = (t_exp, d_exp, params) 
        p_guess = np.array([self.params['m_crystal'], self.params['k_crystal'], 
                            self.params['mu_dashpot']])
        p_opt = fmin(objective, p_guess, args=func_args, xtol=10, ftol=10, 
                     maxiter=200, maxfun=50, full_output=1, disp=0, retall=0)

        self.params.update({'m_crystal':p_opt[0][0], 'k_crystal':p_opt[0][1], 
                               'mu_dashpot':p_opt[0][2]})

        SSres = p_opt[1]
        SStot = np.sum((d_exp - np.mean(d_exp))**2)
        R2 = 1 - SSres / SStot
        self.estimates.update({'R2' : R2})

class ResSim(Sim):
    """
    Class to calculate the dynamics (Force, displacement) of a sphere striking
    a planar surface attached to a sping that is anchored to an infinite
    modulus half-space, such as the case for a sphere or tablet striking a
    dynamic load cell.  The physics of the resonance are determined by the
    spring (load cell crystal) parameters: m_crystal, k_crystal, mu_dashpot
    """
    def __init__(self):
        self.params = {
                'h':None,         # drop height (m)
                'R':None,        # sphere radius (m)
                'Es':None,         # sphere modulus (Pa)
                'Ep':None,         # plane modulus (Pa)
                'vs':None,          # sphere Poisson's ratio
                'vp':None,          # plane Poisson's ratio
                'm':None,       # sphere mass (kg)
                'COR':None,          # coefficient of restitution
                'm_crystal':None,    # mass of the crystal/compact/load cell (kg)
                'k_crystal':None,  # stiffness of the spring (N/m)
                'mu_dashpot':None,   # damping factor of the crystal (N s/m)
                        }

    def set_default_params(self):
        """
        Set default parameter values for debugging
        """
        default_params = {
                'h':4.7e-2,       
                'R':4.76e-3,      
                'Es':3.0e9,       
                'Ep':3.0e9,       
                'vs':0.33,        
                'vp':0.33,        
                'm':0.407e-3,     
                'COR':0.8,
                'm_crystal':0.1,  
                'k_crystal':1.0e9,
                'mu_dashpot':3e3,
                }
        self.set_parameters(default_params)

    def Hookean_force(self,x,k):
        """
        Force of spring following Hooke's law
        x - extension/compression from equilibrium length (m)
        k - spring constant (N/m)
        return force (N)
        """
        F = k * x
        return F

    def dashpot_force(self,x,x_prime,mu):
        """ 
        Dissipative force of the Hookean spring-dashpot model
        x - extension/compression of the Hookean spring (m)
        x_prime - rate of extension/compression of the Hookean spring (m/s)
        mu - proportional constant for disspative force (N s/m)
        return force (N)
        """
        F = mu * x_prime
        return F

    def derivatives(self,y,t,m_s,K,eta,m_c,k,mu):
        """
        function to calculate the derivatives
        y - dependent variable of coupled ODEs -
            [deformation depth, rate change deformation depth, 
             spring compression, rate change spring compression] 
            - [m,m/s,m,m/s]
        t - discrete time array for numerical integration (s)
        m_s - mass of the sphere (kg)
        m_c - mass of the crystal (kg)
        k - Hookean constant (N / m)
        K - Hertz_constant ( kg/(m^0.5*s^2) )
        eta - dissipative constant for Hertz ( kg*m^(1/4)/s )
        mu - dissipative constant for dashpot
        return [2x1] numpy array derivatives
        """
        d0 = y[0]
        d1 = y[1]
        D0 = y[2]
        D1 = y[3]
        try:
            D0_prime = D1
            D1_prime = -1./m_c * (-1 * self.Hertz_force(d0,K) 
                     + self.Hookean_force(D0,k) + self.dashpot_force(D0,D1,mu))
            d0_prime = d1
            d1_prime = -1./m_s * (self.Hertz_force(d0,K) 
                     + self.dissipative_force(eta,d0,d1) - m_s * 9.8 
                     + m_s * D1_prime)
            derivative = np.array([d0_prime,d1_prime,D0_prime,D1_prime])
        except (SystemExit, KeyboardInterrupt):
            raise
        return derivative 
 
    def solve(self):
        """
        numerically integrate the differential equation
        return [[nx1],[nxk],[nx1]] 
            [[time (s)], 
            [deformation (m), rate of deformation (m/s), crystal deflection (m), rate of crystal deflection (m/s)], 
            [force (N)]]
        """
        m = self.params['m']
        R = self.params['R']
        h = self.params['h']
        Es = self.params['Es']
        Ep = self.params['Ep']
        vs = self.params['vs']
        vp = self.params['vp']
        COR = self.params['COR']
        m_c = self.params['m_crystal']
        k_c = self.params['k_crystal']
        K = self.Hertz_constant(Es,Ep,vs,vp,R)
        eta = self.dissipative_constant(m,K,COR)
        mu = self.params['mu_dashpot']
        vo = g*(2*h/g)**0.5
        y_init = [0,1.0*vo,0,0.0*vo]
        t = np.linspace(0,10**-3,10**5)
        y = odeint(self.derivatives,y_init,t,(m,K,eta,m_c,k_c,mu))
        mask = ~np.isnan(y[:,0])
        pts = np.sum(mask)
        t = t[mask]
        y = y[mask]
        t = np.array([t[i] for i in range(0,pts,100)])
        y = np.array([y[i,:] for i in range(0,pts,100)])
        f = self.Hertz_force(y[:,0],K) 
        mask = ~np.isnan(f)
        t = t[mask]
        y = y[mask]
        f = f[mask]
        return [t,y,f]
